#include "hilomodule.h"

// -----------------------------------------------------

static PyObject *hiloIOError = PyErr_NewException("hilo.IOError", NULL, NULL);
static PyObject *hiloOperationError = PyErr_NewException("hilo.OperationError", NULL, NULL);

#define RAISE_EXC__( type, msg ) PyErr_SetString( type, msg ); return NULL
#define RAISE_EXC_IO( msg ) RAISE_EXC__( hiloIOError, msg )
#define RAISE_EXC_OP( msg ) RAISE_EXC__( hiloOperationError, msg )

// returns a dictionary of process where each key is a process name and the value is the process Id (use openProcess to get a handle from Id)
//		if RET_PROCESSES_DICT is defined
// otherwise: returns a list of tuples where the first element is the image name and the second element is the process id
extern "C" static PyObject* hilo_returnProcesses( PyObject *self, PyObject *args ) {
#ifdef RET_PROCESSES_DICT
	PyObject *retDict = PyDict_New();
	
	PROCESSENTRY32 processEntry;
	HANDLE hSnapshot = CreateToolhelp32Snapshot ( TH32CS_SNAPPROCESS, NULL );
	processEntry.dwSize = sizeof(processEntry);

	if ( hSnapshot && Process32First( hSnapshot, &processEntry ) ) {
		PyDict_SetItemString( retDict, processEntry.szExeFile, Py_BuildValue("I", processEntry.th32ProcessID) );

		while ( Process32Next( hSnapshot, &processEntry ) ) {
			PyDict_SetItemString( retDict, processEntry.szExeFile, Py_BuildValue("I", processEntry.th32ProcessID) );
		}
	}

	return retDict;
#else
	PyObject *retList = PyList_New( 0 );

	PROCESSENTRY32 processEntry;
	HANDLE hSnapshot = CreateToolhelp32Snapshot ( TH32CS_SNAPPROCESS, NULL );
	processEntry.dwSize = sizeof(processEntry);

	if ( hSnapshot && Process32First( hSnapshot, &processEntry ) ) {
		PyList_Append( retList, Py_BuildValue( "(sI)", processEntry.szExeFile, processEntry.th32ProcessID ) );

		while ( Process32Next( hSnapshot, &processEntry ) ) {
			PyList_Append( retList, Py_BuildValue( "(sI)", processEntry.szExeFile, processEntry.th32ProcessID ) );
		}
	}

	return retList;
#endif
}

// reads an integer of the specified length and returns its value
// returns a signed value ( could likely change with a parameterization
// exceptions on failure
extern "C" static PyObject* hilo_readInt( PyObject *self, PyObject *args ) {
	// arguments: hilo_HANDLE object, address, length

	hilo_HANDLE *hProc;
	DWORD address, length, buf = 0;

	if ( !PyArg_ParseTuple(args,"OII", &hProc, &address, &length ) ) {
		RAISE_EXC_IO( "Couldn't parse arguments to hilo_readInt" );
	}

	if( readMemory( hProc->hHandle, address, length, &buf ) == NULL ) {
		RAISE_EXC_OP( "Couldn't read process memory" );
	}

	return Py_BuildValue( "i", buf );
}

// writes a value to the memory space of the process given its handle
// note: takes a signed value
// exceptions on failure
// returns non-zero on success
extern "C" static PyObject *hilo_writeInt( PyObject *Self, PyObject *args ) {
	// arguments: hilo_HANDLE object, address, n, val
	hilo_HANDLE *hProc;
	DWORD address, length, val;

	if ( !PyArg_ParseTuple( args, "OIIi", &hProc, &address, &length, &val ) ) {
		RAISE_EXC_IO( "Couldn't parse arguments to hilo_writeInt" );
	}

	if ( writeMemory( hProc->hHandle, address, length, &val ) == NULL ) {
		RAISE_EXC_OP( "Failed on write to memory" );
	}

	return Py_BuildValue("i", 1);
}

// returns the process Id given some name to search for
// uses getPid(), which prioritizes executable name first and window name last
extern "C" static PyObject * hilo_getPid( PyObject *self, PyObject *args) {
	// arguments: name to search for
	const char *name;
	DWORD pid = NULL;

	if (!PyArg_ParseTuple(args, "s", &name)) {
		RAISE_EXC_IO( "Couldn't parse arguments to hilo_getPid" );
	}

	pid = getPid( name );
	return Py_BuildValue( "i", pid );
}

// opens a process by its process Id and returns the handle
// for now, opens it with all relevant desird accesses
// exceptions on failure
// expects the calling function to close the handle before deallocation
extern "C" static PyObject *hilo_openProcess( PyObject *self, PyObject *args ) {
	// arguments: pid
	DWORD pid;
	HANDLE hProc;

	if ( !PyArg_ParseTuple( args, "I", &pid ) ) {
		RAISE_EXC_IO( "Couldn't parse arguments to hilo_openProcess" );
	}

	hProc = openProcess( pid );

	if ( hProc == NULL ) {
		RAISE_EXC_OP( "Failed operation hilo_openProcess" );
	}

	return hilo_HINT::fromHandle( &hProc );
}

// makes a remote thread in the process given by the handle and returns a handle to that thread
// takes a second parameter integer, which is high if the thread should start with TRAP_FLAG set
//		else starts the thread with TRAP_FLAG off
// exceptions if the thread could not be made
// expects the calling function to close the handle before deallocation
extern "C" static PyObject * hilo_makeRemoteThread( PyObject *self, PyObject *args ) {
	// arguments: hilo_HANDLE, int TF

	HANDLE		hThread;	// handle to the created thread
	CONTEXT		cContext;	// thread CONTEXT
	hilo_HANDLE *hProc;		// hilo_HANDLE to the remote process
	DWORD		tf_value;	// 1 or 0 depending on tf

	cContext.ContextFlags = CONTEXT_INTEGER | CONTEXT_CONTROL;

	if ( !PyArg_ParseTuple( args, "OI", &hProc, &tf_value ) ) {
		RAISE_EXC_IO( "Couldn't parse arguments hilo_makeRemoteThread" );
	}

	hThread = makeRemoteSuspendedThread( hProc->hHandle );
	
	if ( hThread == NULL ) {
		RAISE_EXC_OP( "Failure making a remote thread" );
	}

	if ( tf_value ) {
		DEBUG_PRINT( "Starting a remote thread with TRAP_FLAG set" );

		if ( GetThreadContext( hThread, &cContext ) == NULL ) {
			DEBUG_PRINT_VAR( GetLastError() );
			RAISE_EXC_OP( "Could not retrieve thread context" );
		}

		cContext.EFlags |= 1 << 8;

		if ( SetThreadContext( hThread, &cContext ) == NULL ) {
			DEBUG_PRINT_VAR( GetLastError() );
			RAISE_EXC_OP( "Could not set thread context" );
		}
	}

	return hilo_HINT::fromHandle( &hThread );
}

// resumes the thread given by the hilo_HANDLE object and waits for it to end
// if given a second handle, flushes the instruction cache
// does not close the handle to the thread
extern "C" static PyObject * hilo_resumeAndWait( PyObject *self, PyObject *args ) {
	// arguments: hilo_HANDLE thread, opt hilo_HANDLE proc
	hilo_HANDLE *hThread, *hProc = NULL;

	if ( !PyArg_ParseTuple( args, "O|O", &hThread, &hProc ) ) {
		RAISE_EXC_IO( "Couldn't parse arguments to hilo_resumeAndWait" );	
	}

	if ( hProc != NULL ) {
		DWORD buf;
		CONTEXT cContext;
		cContext.ContextFlags = CONTEXT_CONTROL;
		GetThreadContext( hThread->hHandle, &cContext );
		FlushInstructionCache( hProc->hHandle, (LPCVOID)cContext.Eip, 10 );
		VirtualProtectEx( hProc->hHandle,
			(LPVOID)cContext.Eip,
			10,
			PAGE_EXECUTE_READWRITE,
			&buf
			);
	}

	resumeAndWait( hThread->hHandle, false );

	return Py_BuildValue( "i", 1 );
}

// returns the context of the thread given by the HANDLE
// exceptions on failure
extern "C" static PyObject * hilo_getThreadContext( PyObject *self, PyObject *args ) {
	// arguments: hilo_HANDLE (thread handle)
	CONTEXT cContext;
	cContext.ContextFlags = CONTEXT_INTEGER | CONTEXT_CONTROL;
	hilo_HANDLE *hThread;

	if ( !PyArg_ParseTuple( args, "O", &hThread ) ) {
		RAISE_EXC_IO( "Couldn't parse arguments to hilo_getThreadContext" );
	}

	if ( GetThreadContext( hThread->hHandle, &cContext ) == NULL ) {
		DEBUG_PRINT("%s=%d", "GetLastError()=",GetLastError());
		RAISE_EXC_OP( "Failure on retrieving thread context" );
	}

	return hilo_CINT::fromContext( &cContext );
}

// sets the context of the thread given by the HANDLE and the CONTEXT object
// exceptions on failure
extern "C" static PyObject *hilo_setThreadContext( PyObject *self, PyObject *args ) {
	// arguments: hilo_HANDLE, hilo_CONTEXT
	hilo_HANDLE *hThread;
	hilo_CONTEXT *cContext;

	if ( !PyArg_ParseTuple( args, "OO", &hThread, &cContext ) ) {
		RAISE_EXC_IO( "Couldn't parse arguments to hilo_setThreadContext" );
	}

	//cContext->cContext.ContextFlags = CONTEXT_INTEGER | CONTEXT_CONTROL;

	if ( SetThreadContext( hThread->hHandle, &(cContext->cContext) ) == NULL ) {
		DEBUG_PRINT( "%s=%d", "GetLastError()=",GetLastError() );
		RAISE_EXC_OP( "Failure on setting a thread context" );
	}

	return Py_BuildValue( "i", 1 );
}

// closes the handle given
extern "C" static PyObject* hilo_closeHandle( PyObject *self, PyObject *args ) {
	hilo_HANDLE *hHandle;

	if ( !PyArg_ParseTuple( args, "O", &hHandle ) ) {
		RAISE_EXC_IO( "Couldn't parse arguments to hilo_closeHandle" );
	}

	CloseHandle( hHandle );
	return Py_BuildValue( "i", 1 );
}

// returns the default return address of a thread (unsigned integer)
extern "C" static PyObject* hilo_getDefaultReturnAddress( PyObject *self, PyObject *args ) {
	return Py_BuildValue( "I", ExitThread );
}

// returns non-zero if the process given by the handle is being debugged
extern "C" static PyObject* hilo_isProcessBeingDebugged( PyObject *self, PyObject *args ) {
	// arguments: hilo_HANDLE (process)
	hilo_HANDLE *hProc;
	BOOL bRes;

	if ( !PyArg_ParseTuple( args, "O", &hProc ) ) {
		RAISE_EXC_IO( "Couldn't parse arguments to hilo_isProcessBeingDebugged" );
	}

	if ( CheckRemoteDebuggerPresent( hProc->hHandle, &bRes ) == NULL ) {
		DEBUG_PRINT( "%s=%d","GetLastError()",GetLastError() );
		RAISE_EXC_OP( "Could not check if target process was being debugged" );
	}
	return Py_BuildValue( "I", bRes );
}

// creates and attaches a thread to be the debugger of the target process given by the process' handle
// exceptions when fails (such as the process is already being debugged)
extern "C" static PyObject* hilo_traceIntoCall( PyObject *self, PyObject *args ) {
	// arguments: hilo_HANDLE (process), callback func, (remote) fn, args to fn, address
	hilo_HANDLE *_hProc;
	PyObject *callback, *fn, *fnArgs; 
	DWORD threadExitCode = STILL_ACTIVE, debugAddress = NULL;

	// parse arguments
	if ( !PyArg_ParseTuple( args, "OOOO|I", &_hProc, &callback, &fn, &fnArgs, &debugAddress ) ) {
		RAISE_EXC_IO( "Could not parse arguments to hilo_traceCall" );
	}

	// create a thread to debug if able (exception if debugger present)
	BOOL debuggerPresent;
	CheckRemoteDebuggerPresent( _hProc->hHandle, &debuggerPresent );
	if ( debuggerPresent ) {
		RAISE_EXC_OP( "Could not attach as a debugger to the target process" );
	}

	// create debug setup info on the heap and pass it into a newly created thread
	hilo_DebugInfo *hiloDbgSetup = new hilo_DebugInfo( _hProc, callback, debugAddress );
	HANDLE hThread = hiloDbgSetup->CreateDebuggingThread( );
	if ( hThread == NULL ) { // check for an error in creating the thread
		RAISE_EXC_OP( "Could not create debugger thread" );
		DEBUG_PRINT_VAR( GetLastError() );
		delete hiloDbgSetup;
	}

	// synchronize so that a BP is placed before a remote call is made
	DWORD returnState = WaitForSingleObject( hiloDbgSetup->hReadyState, 10000 );

	if ( returnState == WAIT_TIMEOUT || returnState == WAIT_FAILED ) {
		RAISE_EXC_OP( "Wait on synch-semaphore failed" );
		// TODO: clean the thread the we were waiting on
	}

	DEBUG_PRINT("before remote call; after wait");

	if ( debugAddress == NULL ) {
		// call the remote function
		PyObject *fnArgsAsTuple = Py_BuildValue( "(OI)", fnArgs, debugAddress == NULL );
		Py_XINCREF( fnArgsAsTuple );
		PyObject_CallObject( fn, fnArgsAsTuple );
		Py_XDECREF( fnArgsAsTuple );
	} else {
		WaitForSingleObject( hThread, WAIT_TRACE_TIMEOUT );
	}

	// determine if the debugging thread is still active
	GetExitCodeThread( hThread, &threadExitCode );
	if ( threadExitCode == STILL_ACTIVE ) { // if active, report exception
		DEBUG_PRINT( "Badly terminated debugging thread; Waiting for timeout amount (%d ms)", DEBUG_THREAD_TIMEOUT );
		// signal thread back via semaphore and wait
		ReleaseSemaphore( hiloDbgSetup->hReadyState, 1, NULL );
		WaitForSingleObject( hThread, INFINITE );
		//Sleep( DEBUG_THREAD_TIMEOUT );
	}

	CloseHandle( hThread );

	DEBUG_PRINT( "Finished trace" );

	// return a high value
	return Py_BuildValue( "i", 1 );
}

// allocates the specified amount of memory in the target process with the access PAGE_EXECUTE_READWRITE
// returns the allocated address
extern "C" static PyObject* hilo_virtualAlloc( PyObject *self, PyObject *args ) {
	// arguments: hilo_HANDLE (proc), n amount, 
	hilo_HANDLE *hProc;
	DWORD bytes, address;

	if ( !PyArg_ParseTuple( args, "OI", &hProc, &bytes ) ) {
		RAISE_EXC_IO( "Could not parse args to hilo_virtualAlloc" );
	}

	address = (DWORD)VirtualAllocEx(
		hProc->hHandle,
		NULL,
		bytes,
		MEM_COMMIT,
		PAGE_EXECUTE_READWRITE
		);

	if ( address == NULL ) {
		RAISE_EXC_OP( "Unable to allocate memory in the target process." );
		DEBUG_PRINT_VAR( GetLastError() );
	}

	return Py_BuildValue("I",address);
}

// frees some amount of memory starting at the given address in the process specified by the given handle
// returns on success
extern "C" static PyObject* hilo_virtualFree( PyObject *self, PyObject *args ) {
	// arguments: hilo_HANDLE (proc), n amount, address
	hilo_HANDLE *hProc;
	DWORD bytes, address;
	BOOL success;

	if ( !PyArg_ParseTuple( args, "OII", &hProc, &bytes, &address ) ) {
		RAISE_EXC_IO( "Could not parse args to hilo_virtualFree" );
	}

	success = VirtualFreeEx( 
		hProc->hHandle,
		(LPVOID)address,
		bytes,
		MEM_DECOMMIT
		);

	if ( success == NULL ) {
		RAISE_EXC_OP( "Unable to free memory in the target process." );
		DEBUG_PRINT_VAR( GetLastError() );
	}

	return Py_BuildValue("I",success);
}

// reads some memory of arbitrary length at the given address into a buffer and returns it as a ByteArray object
// limit on length: BYTE_ARRAY_BUF_SIZE
// returns on success
extern "C" static PyObject* hilo_readMemory( PyObject *self, PyObject *args ) {
	hilo_HANDLE *_hProc;
	PyObject *byteArray;
	DWORD dwAddress, dwLength;
	char buf[BYTE_ARRAY_BUF_SIZE];

	if ( !PyArg_ParseTuple( args, "OII", &_hProc, &dwAddress, &dwLength ) ) {
		RAISE_EXC_IO( "Could not parse args to hilo_readMemory" );
	}

	if ( dwLength > BYTE_ARRAY_BUF_SIZE ) {
		RAISE_EXC_OP( "In hilo_readMemory, length to read was greater than supported" );
	}

	// set to 0s
	memset( buf, 0, BYTE_ARRAY_BUF_SIZE );

	if ( readMemory( _hProc->hHandle, dwAddress, dwLength, buf ) == NULL ) {
		RAISE_EXC_OP( "Could not read from memory in hilo_readMemory" );
	}

	byteArray = PyByteArray_FromStringAndSize( buf, dwLength );
	return byteArray;
}

// writes some memory of arbitrary length to the given address from a ByteArray object
// returns on success
extern "C" static PyObject* hilo_writeMemory( PyObject *self, PyObject *args ) {
	hilo_HANDLE *_hProc;
	PyObject *byteArray;
	DWORD dwAddress, dwLength;
	char *buf;

	if ( !PyArg_ParseTuple( args, "OIIO", &_hProc, &dwAddress, &dwLength, &byteArray ) ) {
		RAISE_EXC_IO( "Could not parse args to hilo_writeMemory" );
	}

	Py_XINCREF( byteArray );

	buf = PyByteArray_AsString( byteArray );

	if ( writeMemory( _hProc->hHandle, dwAddress, dwLength, buf ) == NULL ) {
		RAISE_EXC_OP( "Could not write to memory in hilo_readMemory" );
	}

	Py_XDECREF( byteArray );

	return byteArray;
}

static PyMethodDef HiloMethods[] = {
	{"getPid", hilo_getPid, METH_VARARGS,"Get process Id from executable file/window title"},
	{"readInt", hilo_readInt, METH_VARARGS,"Returns the integer read from memory specified by process handle, address, and length (max of 4 bytes)."},
	{"writeInt", hilo_writeInt, METH_VARARGS,"Writes an integer to memory specified by process handle, address, length (max of 4 bytes), and val."},
	{"getDefaultReturnAddress", hilo_getDefaultReturnAddress, METH_VARARGS,"Returns the default return address of threads."},
	{"closeHandle", hilo_closeHandle, METH_VARARGS,"Closes the given handle."},
	{"getThreadContext", hilo_getThreadContext, METH_VARARGS,"Gets the context of the given thread."},
	{"setThreadContext", hilo_setThreadContext, METH_VARARGS,"Sets the context of the given thread."},
	{"resumeAndWait", hilo_resumeAndWait, METH_VARARGS,"Resumes the given thread and waits for it to exit"},
	{"openProcess", hilo_openProcess, METH_VARARGS,"Opens the process specified by its process Id and returns the handle."},
	{"makeRemoteThread", hilo_makeRemoteThread, METH_VARARGS,"Creates a thread in the process given by the handle and returns a handle to the thread."},
	{"traceIntoCall", hilo_traceIntoCall, METH_VARARGS,"Sets up a thread for debugging the target process and calls the given remote function with the given arguments."},
	{"virtualAlloc", hilo_virtualAlloc, METH_VARARGS,"Allocates an amount of memory in the process specified by handle, returning the address of the allocated memory."},
	{"virtualFree", hilo_virtualFree, METH_VARARGS,"Frees an amount of memory in the process specified by handle, returning on success."},
	{"readMemory", hilo_readMemory, METH_VARARGS,"Reads memory into a buffer and returns a ByteArray object containing the read bytes."},
	{"writeMemory", hilo_writeMemory, METH_VARARGS,"Writes memory from a ByteArray object into the process and address given."},
	{"getProcesses", hilo_returnProcesses, METH_VARARGS, "Returns a dictionary mapping executable names to process Ids."},
	{NULL, NULL, 0, NULL}
};

PyMODINIT_FUNC inithilo( void ) {
	PyObject* m = Py_InitModule("hilo", HiloMethods);

	Py_INCREF( hiloIOError );
	Py_INCREF( hiloOperationError );
	PyModule_AddObject( m, "IOError", hiloIOError );
	PyModule_AddObject( m, "OperationError", hiloOperationError );

	// define Python CONTEXT object type
	hilo_CINT::define( m );
	// define Python HANDLE object type
	hilo_HINT::define( m );
}